import * as tfvis from '@tensorflow/tfjs-vis'
import * as tf from '@tensorflow/tfjs'
import { getIrisData, IRIS_CLASSES } from './data.js'

/**
 * 多分类 鸢尾花问题
 */
window.onload = async () => {
	const [xTrain, yTrain, xTest, yTest] = getIrisData(0.15)

	xTrain.print()
	console.log("xTrain: ", xTrain)
	console.log("yTrain: ", yTrain)
	console.log(xTest)
	console.log(yTest)

	/**
	 * 连续的模型
	 */
	let model = null
	try {
		model = await tf.loadLayersModel('localstorage://my-model')
	} catch (e) {}
	if (!model) {
		model = tf.sequential()

		/** 隐藏全链接层1
		 * units 神经元数
		 * inputShape 输入形状
		 */
		model.add(tf.layers.dense({
			units: 10,
			inputShape: [xTrain.shape[1]], //4个输入,3个输出问题
			activation: 'relu'
		}))

		/** 输出全链接层
		 * units 神经元数
		 * softmax : 使3个输出 总和为1
		 */
		model.add(tf.layers.dense({
			units: 3,
			activation: 'softmax'
		}))

		/**
		 * 设置损失函数: 交叉熵 categoricalCrossentropy 是对数函数的多分类版本
		 * 优化器adam(学习率)
		 * metrics: 度量单位
		 */
		model.compile({
			loss: "categoricalCrossentropy",
			optimizer: tf.train.adam(0.1),
			metrics: ['accuracy']
		})

		/** fit 开始学习任务
		 * batchSize 取样数
		 * epochs 对训练数据数组进行迭代的次数。
		 * loss: 训练集损失
		 * val_loss: 验证集损失
		 * acc: 训练集准确度
		 * val_acc: 验证集准确度
		 */
		model.fit(xTrain, yTrain, {
			batchSize: 40,
			epochs: 100,
			validationData: [xTest, yTest],
			callbacks: tfvis.show.fitCallbacks({
				name: "训练过程"
			}, [
				'loss', 'val_loss', 'acc', 'val_acc'
			], {
				callbacks: ['onEpochEnd'] //onEpochEnd 结束时
			})
		})

		await model.save('localstorage://my-model');
	}

	window.predict = (form) => {
		const input = tf.tensor([[
			form.a.value * 1,
			form.b.value * 1,
			form.c.value * 1,
			form.d.value * 1
		]])

		const output = model.predict(input);
		// argMax(n) 表示第 n 维最大值
		alert(`预测为 ${IRIS_CLASSES[output.argMax(1).dataSync()[0]]}`);
	}
}

